#!/bin/bash
# 将data/train/refimgs中的图片按比例分配到train、val、test目录
# 默认比例: train:val:test = 8:1:1

# 设置目录路径
SOURCE_DIR="data/train/refimgs"
TRAIN_DIR="data/train/refimgs"
VAL_DIR="data/val/refimgs"
TEST_DIR="data/test/refimgs"

# 创建目标目录
mkdir -p "$VAL_DIR"
mkdir -p "$TEST_DIR"

# 检查源目录是否存在
if [ ! -d "$SOURCE_DIR" ]; then
    echo "错误: 源目录 $SOURCE_DIR 不存在!"
    exit 1
fi

# 获取所有图片文件
echo "正在扫描图片文件..."
cd "$SOURCE_DIR"
IMAGE_FILES=(*.jpg *.jpeg *.png *.bmp *.tiff *.tif *.webp)

# 过滤掉不存在的文件模式
FILTERED_FILES=()
for file in "${IMAGE_FILES[@]}"; do
    if [ -f "$file" ]; then
        FILTERED_FILES+=("$file")
    fi
done

# 检查是否找到图片文件
if [ ${#FILTERED_FILES[@]} -eq 0 ]; then
    echo "错误: 在 $SOURCE_DIR 中没有找到图片文件!"
    exit 1
fi

TOTAL_FILES=${#FILTERED_FILES[@]}
echo "找到 $TOTAL_FILES 个图片文件"

# 计算分配数量 (8:1:1比例)
VAL_COUNT=$((TOTAL_FILES / 10))
TEST_COUNT=$((TOTAL_FILES / 10))
TRAIN_COUNT=$((TOTAL_FILES - VAL_COUNT - TEST_COUNT))

echo "分配方案:"
echo "  训练集: $TRAIN_COUNT 个文件 (保留在原目录)"
echo "  验证集: $VAL_COUNT 个文件 (移动到 $VAL_DIR)"
echo "  测试集: $TEST_COUNT 个文件 (移动到 $TEST_DIR)"

# 随机打乱文件列表
shuf_files=($(printf '%s\n' "${FILTERED_FILES[@]}" | shuf))

# 移动文件到验证集
echo "正在移动文件到验证集..."
for ((i=0; i<VAL_COUNT; i++)); do
    file="${shuf_files[i]}"
    mv "$file" "../../val/refimgs/"
    echo "  移动: $file -> val/refimgs/"
done

# 移动文件到测试集
echo "正在移动文件到测试集..."
for ((i=VAL_COUNT; i<VAL_COUNT+TEST_COUNT; i++)); do
    file="${shuf_files[i]}"
    mv "$file" "../../test/refimgs/"
    echo "  移动: $file -> test/refimgs/"
done

echo "数据集分割完成!"
echo "最终分布:"
echo "  训练集: $(ls -1 "$TRAIN_DIR" | wc -l) 个文件"
echo "  验证集: $(ls -1 "$VAL_DIR" | wc -l) 个文件"
echo "  测试集: $(ls -1 "$TEST_DIR" | wc -l) 个文件"

echo ""
echo "接下来可以运行以下命令生成失真数据:"
echo "  python utils/data_generate/live_generator.py --data data/train/refimgs"
echo "  python utils/data_generate/live_generator.py --data data/val/refimgs"
echo "  python utils/data_generate/live_generator.py --data data/test/refimgs"